#include "stdafx.h"

#include <opencv2/legacy/legacy.hpp>
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/photo/photo.hpp"
#include "opencv2/features2d/features2d.hpp"
#include "opencv2/nonfree/features2d.hpp"
#include "opencv2/calib3d/calib3d.hpp"
#include "opencv2/video/background_segm.hpp"

#include "BackgroundSegmentation.h"
#include "OpticalFlow.h"
#include "ImageMethods.h"
#include "ImageStitcher.h"
#include "RGBColour.h"
#include "PatchFinder.h"
#include "ResourceLibrary.h"

using namespace cv;
using namespace std;

Point2f pt;
bool addRemovePt = false;

class BackgroundSubtractorMOG2Custom: public BackgroundSubtractorMOG2{
public:
    void setTau(float val){fTau = val;}
	void setVarInit(float val){fVarInit = val;}
	void setFCT(float val){fCT = val;}
	void setVarThresholdGen(float val){varThresholdGen = val;}
	void setBackgroundRatio(float val){backgroundRatio = val;}
};

BgSegmenter::BgSegmenter(void)
{
	foreground;
	frame;
	image;
	prev_frame;
    mog;
	rl = new ResourceLibrary();
	of = new OpticalFlow(rl);
	im = new ImageMethods();
	is = new ImageStitcher(rl);
	pf = new PatchFinder(rl,is);
	images;
	originals;
}

BgSegmenter::~BgSegmenter(void)
{
}

Mat BgSegmenter::getCurrentFrame()
{
	return foreground;
}

Mat BgSegmenter::getPrevFrame()
{
	return prev_frame;
}

/*
* Function to test that the bg/fg sgmentation images are being stored in images vector
*/
void BgSegmenter::showImages()
{
	namedWindow( "Foreground ", CV_WINDOW_AUTOSIZE );
	for(int i=0;i<images.size();i++)
	{
		cvShowImage("Foreground", &images.at(i));
		waitKey(100);
	}
}

/*
* Function to save the temporal correction images
*/
void BgSegmenter::saveCorrections()
{
	for(int i=0;i<images.size();i++)
	{
		std::stringstream sstm;
		sstm << "images/fgbg/frame" << i;
		String filename = sstm.str();
		im->saveImage(&images.at(i),filename);
	}
}

/*
* Function to remove temporal noise from segmented images
*/
void BgSegmenter::removeTemporal(int fs)
{
	IplImage* frame;
	IplImage* tests;
	RGBColour blackPix = RGBColour(0,0,255);
	unsigned char px; 
	int bg = 0;
	int fg = 0;
	//namedWindow( "Frame", CV_WINDOW_AUTOSIZE );
	/* For all background segmentation images */
	for (int i = fs; i<(images.size()-fs);i++)
	{
		frame = &images.at(i);
		/* For all pixels in the image */
		for(int x = 0; x<frame->width;x++)
		{
			for(int y =0; y<frame->height;y++)
			{
				px = im->getPixel(frame,x,y);
				if(px == 255)
				{
					/* For images either side of the current image */
					for(int n=(i-fs);n<(1+i+fs);n++)
					{
						tests=&images.at(n);
						px = im->getPixel(tests,x,y);
						if(px == 255)
						{fg++;}
						else if (px==0)
						{bg++;}
					}
					if(bg>fg)
					{
						im->setPixel(frame,x,y,0);
						//cvShowImage("Frame", &images.at(i));
						//waitKey(100);
					}
					bg=0;
					fg=0;
				}
			}
		}
		std::stringstream sstm;
		sstm << "bg" << i;
		String filename = sstm.str();
		im->saveImage(frame,filename);
	}
}

/*
*	Background segmentation correction using optical flow information
*/
void BgSegmenter::removeOpticalFlow(int f)
{
	//IplImage* img = &images.at(f);
	Mat img(&images.at(f));
	CvPoint p;
	CvPoint q;
	int thickness = -1;
	int lineType = 8;
	int radius = 10;
	for (int n=0;n<of->goodEnd.size();n++)
	{
		p = of->goodEnd.at(n);
		q = of->goodStart.at(n);
		circle( img, p, radius, Scalar( 0, 0, 255 ), thickness, lineType );
		circle( img, q, radius, Scalar( 0, 0, 255 ), thickness, lineType );
		imshow("circles",img);
	}
}

/*
*	HIghlights pixels that change TO foreground between frames
*/
IplImage* BgSegmenter::subtractOpticalFlow(IplImage* curr_frame, IplImage* fg)
{
	frame = curr_frame;
	if (fg != NULL)
	{
		prev_frame = fg;
		imshow("Previous Frame ",prev_frame);
	}
    image=frame.clone();
    mog(frame,foreground,-1);
    threshold(foreground,foreground,128,255,THRESH_BINARY);
    medianBlur(foreground,foreground,9);
    erode(foreground,foreground,Mat());
    dilate(foreground,foreground,Mat());

    imshow("Current Frame ",foreground);
	if (fg != NULL)
	{
		im->compareGreyImages(prev_frame,foreground);
	}
	IplImage* output=cvCloneImage(&(IplImage)foreground);

	return output;
 }

/*
* Main background segmentation function which calls optical flow calculation
*/
int BgSegmenter::segmentBackground(int filter_size, String filename, int save)
{

	createImages(filename);
	stitchImages();
	//removeTemporal(filter_size);
    if(save==1)
	{
		saveCorrections();
	}
 	//of->calcOpticalFlow(this);
	pf->defineROI(Point(250,300),200,300);
	//pf->displayWarpedImages();
	pf->findValidPatches();
	pf->displayValidPatches();
	pf->savePatches();

	return 0;
}

/*
* Initial function to populate imkages vector with the initial segmentation images
*/
int BgSegmenter::createImages(String filename)
{
	IplImage* storeImage;
	IplImage* storeOrig;
	VideoCapture cap;
    cap.open(filename);
    if( !cap.isOpened() )
    {
        puts("***Could not initialize capturing...***\n");
        return 0;
    }
    //namedWindow( "Capture ", CV_WINDOW_AUTOSIZE);
    namedWindow( "Foreground ", CV_WINDOW_AUTOSIZE );
    Mat frame,foreground,image;
    BackgroundSubtractorMOG2Custom mog;

	mog.setTau(0.5);
	mog.setBackgroundRatio(0.92);
	mog.setVarInit(15);
	mog.setVarThresholdGen(30);
	mog.setFCT(0.5);

    int fps=cap.get(CV_CAP_PROP_FPS);
    if(fps<=0)
        fps=10;
    else
        fps=1000/fps;
    for(;;)
    {
        cap>>frame;   
        if(frame.empty())
		{
			break;
		}
        image=frame.clone();
		//mog.set("backgroundRatio",0.9);
        mog(frame,foreground,-1);
        threshold(foreground,foreground,128,255,THRESH_BINARY);
        medianBlur(foreground,foreground,9);
        erode(foreground,foreground,Mat());
        dilate(foreground,foreground,Mat());

        //imshow( "Capture ",image );
        imshow("Foreground ",foreground);

		rl->addOriginals(frame);
		rl->addSegmentedImages(foreground);
		//storeImage=cvCloneImage(&(IplImage)foreground);
		//images.push_back(*storeImage);
		//storeOrig=cvCloneImage(&(IplImage)frame);
		//originals.push_back(*storeOrig);

        //char c = (char)waitKey(fps);
        //if( c == 27 ) 
           // break;

		//waitKey(0);
    }
	return 0;
}

int BgSegmenter::stitchImages() 
{
	Mat result;
	for(int i=0;i<(rl->originals.size()-1);i=i+1)
	{
		Mat im1(&rl->originals.at(i));
		Mat im2(&rl->originals.at(i+1));
		if (i==0)
		{
			result = is->stitchImages(im2,im1);
		}
		else
		{
 			result = is->stitchImages(im2,result);
		}
		rl->addWarpedImages(result);
	}
	return 0;
}
